#  -*- coding: utf-8 -*-

from strategy.stock_pool.base_stock_pool import BaseStockPool
from factor.factor_module import FactorModule
from data.data_module import DataModule
from util.stock_util import get_trading_dates,get_all_codes_date,get_diff_dates,calc_negative_diff_dates,calc_positive_diff_dates,multi_computer,get_all_codes_trading_date
from pandas import DataFrame,Series
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import pylab as pl
import scipy.signal as signal
import time
from sklearn import linear_model
import pickle
import os
"""
趋势策略股票池定义
A:波段1起点（极小值）；
B：波段1终点（极大值）；
C：波段2回调最低点（极小值）；
D：当前日；
"""


class ShortStrengthStockPool(BaseStockPool):

    def multi_get_short_strength_stock(self,date):
        start_time = time.time()
        all_codes = get_all_codes_trading_date(date)
        args =(date,)
        option_codes_dict = dict()
        option_codes_multi_lst = multi_computer(get_short_strength_stock_at_date,all_codes,args)
        #print(option_codes_multi_lst)
        for codes_dict in option_codes_multi_lst:
            option_codes_dict.update(codes_dict)

        end_time = time.time()
        #print(option_codes_dict)
        print("多进程执行短线梅开二度策略选股，日期：%s,候选股票：%s,耗时：%.2f" %(date,option_codes_dict.keys(),(end_time - start_time)))
        return option_codes_dict

    def pickle_option_stocks(self):
        file_dates_name = "short_strength_stock_pool_rebalance_dates_" + self.begin_date + "_" + self.end_date
        file_codes_name = "short_strength_stock_pool_codes_dict_" + self.begin_date + "_" + self.end_date
        file_dates_path = os.path.join(self.stock_pool_path,file_dates_name)
        file_codes_path = os.path.join(self.stock_pool_path,file_codes_name)
        rebalance_dates, rebalance_date_codes_dict = self.get_option_stocks()

        with open(file_dates_path, 'wb') as f:
            pickle.dump(rebalance_dates, f)
        f.close()
        with open(file_codes_path, 'wb') as f:
            pickle.dump(rebalance_date_codes_dict, f)
        f.close()

    def get_pickle_option_stocks(self):

        codes_dict = dict()
        rebalance_dates = []
        for root, dirs, files in os.walk(self.stock_pool_path):
            for name in files:
                file = os.path.join(root, name)
                if "short_strength_stock_pool_codes_dict" in file:
                    with open(file, 'rb') as f:
                        codes_dict.update(pickle.load(f))
                    f.close()
                if "short_strength_stock_pool_rebalance_date" in file:
                    with open(file, 'rb') as f:
                        rebalance_dates.extend(pickle.load(f))
                    f.close()

        print(rebalance_dates,codes_dict)

        return rebalance_dates,codes_dict

    def get_option_stocks(self):
        """
        实现股票池选股逻辑，找到指定日期范围的候选股票
        条件：0 < PE < 30, 按从小到大排序，剔除停牌后，取前100个；再平衡周期：7个交易日
        :return: tuple，再平衡的日期列表，以及一个dict(key: 再平衡日, value: 当期的股票列表)
        """

        dm = DataModule()

        # 获取日期范围的交易日历
        all_dates = get_trading_dates(self.begin_date, self.end_date)

        # 缓存股票和其对应有交易的日期
        code_dates_cache = dict()

        # 调整日和其对应的股票
        rebalance_date_codes_dict = dict()
        rebalance_dates = []

        # 保存上一期的股票池
        last_phase_codes = []
        # 所有的交易日数
        dates_count = len(all_dates)

        # 用再平衡周期作为步长循环
        for index in range(0, dates_count, self.interval):
            # 当前的调整日
            rebalance_date = all_dates[index]
            print(rebalance_date)
            self.logger.info('调整日：' + rebalance_date)

            # 获取本期符合条件的备选股票
            this_phase_option_codes_dict = self.multi_get_short_strength_stock(rebalance_date)
            #this_phase_option_codes = get_trending_stock_at_date(None,rebalance_date)
            this_phase_option_codes = list(this_phase_option_codes_dict.keys())
            self.logger.info('原始备选股票')
            this_phase_option_codes.sort()
            self.logger.info(this_phase_option_codes)

            # 本期入选的股票代码列表
            this_phase_codes = []

            # 找到在上一期的股票池，但是当前停牌的股票，保留在当期股票池中
            if len(last_phase_codes) > 0:
                print("上一期股票池个数：%4d" % len(last_phase_codes))
                for code in last_phase_codes:
                    if code not in list(code_dates_cache.keys()):
                        daily_ks = dm.get_k_data(code, autype=None,index = False, begin_date=self.begin_date, end_date=self.end_date)
                        daily_ks.set_index(['date'], inplace=True)
                        # 只保留交易日的数据
                        daily_ks = daily_ks[daily_ks['is_trading'] == True]
                        code_dates_cache[code] = list(daily_ks.index)

                    if rebalance_date not in code_dates_cache[code]:
                        this_phase_codes.append(code)


            self.logger.info('上期停牌的股票')
            this_phase_codes.sort()
            self.logger.info(this_phase_codes)

            # 剩余的位置用当前备选股票的,先去重
            for dup_code in this_phase_codes:
                if dup_code in this_phase_option_codes:
                    this_phase_option_codes.remove(dup_code)

            option_size = len(this_phase_option_codes)
            if option_size > (100 - len(this_phase_codes)):
                this_phase_codes += this_phase_option_codes[0:100 - len(this_phase_codes)]

            else:
                this_phase_codes += this_phase_option_codes

            # 当期股票池作为下次循环的上期股票池
            last_phase_codes = this_phase_codes

            # 保存到返回结果中
            rebalance_date_codes_dict[rebalance_date] = this_phase_codes
            rebalance_dates.append(rebalance_date)

            self.logger.info('当前最终的备选票：%s' % rebalance_date)
            this_phase_codes.sort()
            self.logger.info(this_phase_codes)

        return rebalance_dates, rebalance_date_codes_dict

    def generate_feature_data(self):
        # 获取日期范围的交易日历
        all_dates = get_trading_dates(self.begin_date, self.end_date)

        # 记录股票池里股票的特征数据
        code_feature_df = pd.DataFrame(columns=['date','code','date_A','date_B','date_C'])

        for date in all_dates:

            this_phase_option_codes_dict = self.multi_get_short_strength_stock(date)

            for key in this_phase_option_codes_dict.keys():
                code_feature_df = code_feature_df.append(this_phase_option_codes_dict[key],ignore_index=True)

       #print(code_feature_df)s
        file_name = "short_strength_stock_pool_feature_data_" + self.begin_date + "_" + self.end_date
        file_path = os.path.join(self.stock_pool_path,file_name)

        with open(file_path, 'wb') as f:
            pickle.dump(code_feature_df, f)
        f.close()


    def get_feature_data(self):

        code_feature_df = pd.DataFrame(columns=['date','code','date_A','date_B','date_C'])
        for root, dirs, files in os.walk(self.stock_pool_path):
            for name in files:
                file = os.path.join(root, name)
                if "short_strength_stock_pool_feature_data" in file:
                    with open(file, 'rb') as f:
                        code_feature_df = code_feature_df.append(pickle.load(f),ignore_index=True)
                    f.close()
        #去重,根据去重后数目重新从0开始索引
        code_feature_df.drop_duplicates(['code','date'],keep='first',inplace=True)
        code_feature_df['index'] = range(code_feature_df.shape[0])
        code_feature_df.set_index('index',inplace=True)
        return code_feature_df

    def train_data(self):
        code_feature_df = self.get_feature_data()
        code_feature_df["Y"] = 0
        #code_feature_df_copy = code_feature_df.copy()
        dm = DataModule()
        fm = FactorModule()
        for i in code_feature_df.index:
            begin_date = code_feature_df.loc[i]["date"]
            code = code_feature_df.loc[i]['code']
            end_date = calc_positive_diff_dates(code=code, date=begin_date, delta_days=50)
            date_B = code_feature_df.loc[i]["date_B"]
            if end_date == False:
                continue

            code_df = dm.get_k_data(code, index=False, autype="hfq", begin_date=begin_date, end_date=end_date)
            if code_df.index.size == 0:
                continue
            code_df.set_index(['date'], inplace=True)

            #先找整个区间最大值，再找起始日到最大日之间的最小值
            close_ser = code_df["close"]
            beg_min = code_df.loc[begin_date]['low']
            beg_close = code_df.loc[begin_date]['close']
            max_value = close_ser.max()
            max_idx = close_ser.idxmax()
            diff_dates = get_diff_dates(code,begin_date,max_idx)
            #print("maxidx: %s,max_value: %s" %(max_idx,max_value))
            code_to_max_df = code_df.loc[begin_date:max_idx]
            min_value = code_to_max_df['close'].min()
            min_idx = code_to_max_df['close'].idxmin()
            #print("min_idx: %s,min_value: %s" %(min_idx,min_value))

            #最大值日2个交易日（保证可卖出），最大值收益超过8%，最小值不低于起始日的最低价
            max_change_rate = round(100 * (max_value - beg_close)/beg_close,2)
            #print(diff_dates,max_change_rate,min_value,beg_min)
            if diff_dates >= 2 and max_change_rate >= 8  and min_value >= beg_min :
                code_feature_df.loc[i,"Y"] = 1

            date_C_rs_df = fm.get_single_stock_factors(code, "rs", False, begin_date, begin_date)
            date_B_rs_df = fm.get_single_stock_factors(code, "rs", False, date_B, date_B)
            if len(date_C_rs_df) == 1 and "rs_250" in date_C_rs_df.columns.values:
                code_feature_df.loc[i, "date_C_rs_60"] = date_C_rs_df.loc[0]['rs_60']
                code_feature_df.loc[i, "date_C_rs_120"] = date_C_rs_df.loc[0]['rs_120']
                code_feature_df.loc[i, "date_C_rs_250"] = date_C_rs_df.loc[0]['rs_250']

            if len(date_B_rs_df) == 1 and "rs_250" in date_C_rs_df.columns.values:
                code_feature_df.loc[i, "date_B_rs_60"] = date_C_rs_df.loc[0]['rs_60']
                code_feature_df.loc[i, "date_B_rs_120"] = date_C_rs_df.loc[0]['rs_120']
                code_feature_df.loc[i, "date_B_rs_250"] = date_C_rs_df.loc[0]['rs_250']


            #作图
            # code_df.plot(title='close', y=['close'], kind='line')
            # plt.plot(extrem_max_index, close_np[extrem_max_index], 'o')
            # plt.plot(extrem_min_index, close_np[extrem_min_index], '+')
            # for x, y in zip(extrem_max_index, close_np[extrem_max_index]):
            #     plt.text(x, y + 0.3, '%s' % (code_df.index[x]), ha='center', va='bottom', fontsize=10.5)
            # for x, y in zip(extrem_min_index, close_np[extrem_min_index]):
            #     plt.text(x, y + 0.3, '%s' % (code_df.index[x]), ha='center', va='bottom', fontsize=10.5)
            # plt.show()

        print(code_feature_df)


def get_short_strength_stock_at_date(codes,date):
    start_time = time.time()
    option_codes_dict = dict()
    fm = FactorModule()
    if codes == None:
        codes = ['600082']

    dm = DataModule()
    change_rate_minimum = 9.5

    for code in codes:

        try:
            #每次循环需要建一个新的对象，否则赋值后下一个循环会被覆盖掉，从而影响保存的值
            code_data_s = pd.Series(index=['date', 'code', 'date_A', 'date_B', 'date_C'])

            print("开始计算股票：%s的短期强度" % code)
            # 取近3月数据
            begin_date = calc_negative_diff_dates(code=None,date = date,delta_days = -60)
            if begin_date == False:
                begin_date = None

            code_df = dm.get_k_data(code,index=False,autype="hfq",begin_date=begin_date,end_date=date)

            if code_df.index.size == 0:
                continue

            code_df.set_index(['date'], inplace=True)
            ma_20,ma_250,change_rate_above_ma20 = None,None,None

            #第一步，找出当日涨停股，放量
            change_rate_df = fm.get_single_stock_factors(code,"change_rate",False,date,date)
            hfq_ma_df = fm.get_single_stock_factors(code,"hfq_ma",False,date,date)
            vol_ratio_df = fm.get_single_stock_factors(code,"vol_ratio",False,date,date)
            close = code_df.loc[date]['close']
            amount = code_df.loc[date]['amount']
            if len(change_rate_df) == 1:
                change_rate = change_rate_df.loc[0]['close_change_rate']
                price_limit = change_rate_df.loc[0]['price_limit']
            if len(hfq_ma_df) == 1:
                ma_20 = hfq_ma_df.loc[0]['ma20']
                ma_250 = hfq_ma_df.loc[0]['ma250']
                if pd.isnull(ma_20) or pd.isnull(ma_250):
                    print("S1:没有20日线或250日线，为新股或次新股")
                    #continue
                #change_rate_above_ma20 = round(100*(close-ma_20)/ma_20,2)
            if len(vol_ratio_df) == 1:
                vol_ratio_5 = vol_ratio_df.loc[0]["vol_ratio_5"]
                vol_ratio_10 = vol_ratio_df.loc[0]["vol_ratio_10"]
                vol_ma5 = vol_ratio_df.loc[0]["vol_ma5"]
                vol_ma10 = vol_ratio_df.loc[0]["vol_ma10"]
                if pd.isnull(vol_ratio_5) or pd.isnull(vol_ratio_10):
                    print("S1:没有5日或10日量比")
                    continue

            #涨跌幅四舍五入
            print("S1:当日涨幅：%.2f,收盘价：%.2f,成交额：%.2f 亿元,5日量比：%.2f，10日量比：%.2f"
                  % (change_rate,close,amount/1E8,vol_ratio_5,vol_ratio_10))

            if change_rate < change_rate_minimum\
                or amount < 1E8:
                continue
            # if round(change_rate,1) < change_rate_minimum\
            #     or close < ma_20\
            #     or close < ma_250\
            #     or change_rate_above_ma20 > 10\
            #     or amount < 1E8\
            #     or vol_ratio_5 < 1\
            #     or vol_ratio_10 < 1:
            #     continue

            #第二步，找出最近的极大值极小值点,计算波段1和波段2涨跌幅，
            #波段1涨幅>30,波段2回调<30


            extrem_date_index = []
            extrem_date_values = []
            close_array = code_df["close"]
            code_df['ma5'] = code_df["close"].rolling(5).mean()
            close_np = close_array.values

            extrem_max_index = signal.argrelextrema(close_np, np.greater, order=1)[0]
            extrem_min_index = signal.argrelextrema(-close_np, np.greater, order=1)[0]

            extrem_max_date_index = code_df.index[extrem_max_index]
            extrem_min_date_index = code_df.index[extrem_min_index]

            #波段2最大回调10天，所以波段2期间大于10天的极大值点舍去
            [extrem_date_index.append(x) for x in extrem_max_date_index]
            [extrem_date_index.append(x) for x in extrem_min_date_index]
            #降序排序,值的顺序与排序后的日期顺序相同
            extrem_date_index.sort(reverse = True)
            [extrem_date_values.append(code_df.loc[x]['close']) for x in extrem_date_index]

            #print(extrem_date_index)
            date_A,date_A_value = None,None
            date_B,date_B_value = None,None
            date_C,date_C_value = None,None
            for extrem_date in extrem_date_index:
                #找到第一个10天以内的间隔的极大值，在此范围内找最大值，那么就是B
                if date_B == None and extrem_date in extrem_max_date_index:
                    diff_dates = get_diff_dates(code,extrem_date,date)
                    if diff_dates < 10:
                        #list切片是左闭右开，所以要+1
                        date_max_value_most_10 = max(extrem_date_values[0:extrem_date_index.index(extrem_date)+1])
                        date_max_most_10 = extrem_date_index[extrem_date_values.index(date_max_value_most_10)]

                        date_B = date_max_most_10
                        date_B_value = extrem_date_values[extrem_date_index.index(date_B)]

                #找到B之后，B之前（倒序排列意义上的之前，实际意义上的之后）的最小值是C（如果有多个相同的最小值，取最早的一个）
                if date_B is not None:
                    date_C_value = code_df.loc[date_B:date,'close'].min()
                    #date_C_value = min(extrem_date_values[0:extrem_date_index.index(date_B)+1])
                    date_C = code_df[(code_df.index >date_B) & (code_df["close"] == date_C_value)].index.tolist()[0]

                    #找到B之后，B日5天前的交易日为A
                    if date_A == None:
                        for min_date in extrem_date_index[extrem_date_index.index(date_B):]:
                            #print(min_date)
                            if min_date in extrem_min_date_index:
                                date_A = min_date
                                date_A_value = extrem_date_values[extrem_date_index.index(date_A)]
                        #date_A = calc_negative_diff_dates(code=code, date=date_B, delta_days=-6)
                        #date_A_value = code_df.loc[date_A]['open']
                                break


            #A是在找到B的前提下找到的，所以A不为空，其他都不是空
            if date_A is None:
                print("S2:没找到date_A,不符合短期强度形态")
                continue
            #
            change_rate_AB,diff_date_AB,change_rate_BC,diff_date_BC = 0,0,0,0
            change_rate_AB = round(100*(date_B_value - date_A_value)/date_A_value,2)
            diff_date_AB = get_diff_dates(code,date_A,date_B)
            change_rate_BC = round(100 * (date_C_value - date_B_value) / date_B_value, 2)
            diff_date_BC = get_diff_dates(code, date_B, date_C)
            change_rate_BD = round(100*(close - date_B_value) / date_B_value, 2)
            diff_date_BD = get_diff_dates(code,date_B,date)
            print("S2:A日：%s,B日：%s,C日：%s,AB涨幅：%.2f,AB间隔：%s天，BC回调：%.2f，BC间隔：%s天，BD涨幅: %.2f,BD间隔：%.2f"
                   % (date_A,date_B,date_C,change_rate_AB,diff_date_AB,change_rate_BC,diff_date_BC,change_rate_BD,diff_date_BD))

            # clf = linear_model.LinearRegression()
            # date_near_4 = calc_negative_diff_dates(code,False,date,-4)
            # #print(date_near_4)
            # clf.fit(np.arange(5).reshape(-1,1),code_df['ma5'].loc[date_near_4:date])
            # k_20 = clf.coef_
            # print("S2:最近5日的20日线拟合的斜率：%.2f" % clf.coef_)
            # clf.fit(np.arange(get_diff_dates(code,date_C,date)).reshape(-1,1),code_df['close'].loc[date_C:date])
            # k_CD = clf.coef_
            # print("S2:从C点到当日日线收盘价拟合的斜率：%.2f" % clf.coef_)

            #

            if change_rate_AB < 30 or change_rate_BC < -30 or change_rate_BC > 0 \
                or close > date_B_value:
                print("S2:形态不满足量化条件")
                continue

            #第三步，成交量判断
            #波段1成交量大，波段2成交量小
            amount_mean_AB= code_df.loc[date_A:date_B,"amount"].mean()
            amount_mean_BD= code_df.loc[date_B:date,"amount"].mean()
            change_rate_amount_AB_BC = round(100 * (amount_mean_AB - amount_mean_BD)/amount_mean_AB,2)
            print("S3:波段1(AB)平均成交额：%.2f亿元,波段2(BD)平均成交额：%.2f亿元,差值比例: %.2f"
                  %(amount_mean_AB/1E8,amount_mean_BD/1E8,change_rate_amount_AB_BC))
            if amount_mean_AB < amount_mean_BD :
                print("S3: AB量能比BD低，不满足波段1放量，波段2缩量的条件")
                continue
            code_data_s["date"] = date
            code_data_s["code"] = code
            code_data_s["date_A"] = date_A
            code_data_s["date_B"] = date_B
            code_data_s["date_C"] = date_C

            option_codes_dict[code] = code_data_s

            #作图
            code_df.plot(title='close', y=['close',"ma5"], kind='line')
            plt.plot(extrem_max_index, close_np[extrem_max_index], 'o')
            plt.plot(extrem_min_index, close_np[extrem_min_index], '+')
            for x, y in zip(extrem_max_index, close_np[extrem_max_index]):
                plt.text(x, y + 0.3, '%s' % (code_df.index[x]), ha='center', va='bottom', fontsize=10.5)
            for x, y in zip(extrem_min_index, close_np[extrem_min_index]):
                plt.text(x, y + 0.3, '%s' % (code_df.index[x]), ha='center', va='bottom', fontsize=10.5)
            plt.show()
        except:
            print("股票 %s,计算短期强度标是错误，日期：%s" % (code,date))

    end_time = time.time()
    print("日期：%s,候选股票：%s,耗时：%.2f"% (date,option_codes_dict.keys(),(end_time - start_time)))
    return option_codes_dict




if __name__ == '__main__':
    tp = ShortStrengthStockPool(begin_date="2015-01-01",end_date="2018-11-26",interval=1)
    #tp.generate_feature_data()
    #tp.train_data()
    tp.multi_get_short_strength_stock("2019-03-04")
    #print(tp.get_option_stocks())
    #result = get_short_strength_stock_at_date(codes = None,date="2019-03-04")
    #print(result[1])
    #tp.pickle_option_stocks()
    #a,b = tp.get_pickle_option_stocks()

